{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 猫狗识别-数据增强篇\n",
    "\n",
    "阅读本文你将学会**如何通过 PyTorch 在训练模型前进行数据增强**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor, Lambda, Compose\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据读取与预处理（重点内容）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_datadir = './data/cat-dog/'\n",
    "test_datadir  = './data/cat-dog/'\n",
    "\n",
    "train_transforms = transforms.Compose([\n",
    "    transforms.Resize([224, 224]),  # 将输入图片resize成统一尺寸\n",
    "    transforms.RandomRotation(degrees=(-10, 10)),  #随机旋转，-10到10度之间随机选\n",
    "    transforms.RandomHorizontalFlip(p=0.5),  #随机水平翻转 选择一个概率概率\n",
    "    transforms.RandomVerticalFlip(p=0.5),  #随机垂直翻转\n",
    "    transforms.RandomPerspective(distortion_scale=0.6, p=1.0), # 随机视角\n",
    "    transforms.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5)),  #随机选择的高斯模糊模糊图像\n",
    "    transforms.ToTensor(),          # 将PIL Image或numpy.ndarray转换为tensor，并归一化到[0,1]之间\n",
    "    transforms.Normalize(           # 标准化处理-->转换为标准正太分布（高斯分布），使模型更容易收敛\n",
    "        mean=[0.485, 0.456, 0.406], \n",
    "        std=[0.229, 0.224, 0.225])  # 其中 mean=[0.485,0.456,0.406]与std=[0.229,0.224,0.225] 从数据集中随机抽样计算得到的。\n",
    "])\n",
    "\n",
    "test_transforms = transforms.Compose([\n",
    "    transforms.Resize([224, 224]),  # 将输入图片resize成统一尺寸\n",
    "    transforms.ToTensor(),          # 将PIL Image或numpy.ndarray转换为tensor，并归一化到[0,1]之间\n",
    "    transforms.Normalize(           # 标准化处理-->转换为标准正太分布（高斯分布），使模型更容易收敛\n",
    "        mean=[0.485, 0.456, 0.406], \n",
    "        std=[0.229, 0.224, 0.225])  # 其中 mean=[0.485,0.456,0.406]与std=[0.229,0.224,0.225] 从数据集中随机抽样计算得到的。\n",
    "])\n",
    "\n",
    "train_data = datasets.ImageFolder(train_datadir,transform=train_transforms)\n",
    "test_data  = datasets.ImageFolder(test_datadir,transform=test_transforms)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_data,\n",
    "                                          batch_size=4,\n",
    "                                          shuffle=True,\n",
    "                                          num_workers=1)\n",
    "test_loader  = torch.utils.data.DataLoader(test_data,\n",
    "                                          batch_size=4,\n",
    "                                          shuffle=True,\n",
    "                                          num_workers=1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "关于 `transforms.Compose` 这部分更多的信息可以参考 https://pytorch-cn.readthedocs.io/zh/latest/torchvision/torchvision-transform/\n",
    "\n",
    "如果你想知道还有哪些数据增强手段，可以看看这里：https://pytorch.org/vision/stable/auto_examples/plot_transforms.html#sphx-glr-auto-examples-plot-transforms-py 对应的API，你可以在这里找到 https://pytorch.org/vision/stable/transforms.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of X [N, C, H, W]:  torch.Size([4, 3, 224, 224])\n",
      "Shape of y:  torch.Size([4]) torch.int64\n"
     ]
    }
   ],
   "source": [
    "for X, y in test_loader:\n",
    "    print(\"Shape of X [N, C, H, W]: \", X.shape)\n",
    "    print(\"Shape of y: \", y.shape, y.dtype)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device\n",
      "NeuralNetwork(\n",
      "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
      "  (linear_relu_stack): Sequential(\n",
      "    (0): Linear(in_features=150528, out_features=512, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=512, out_features=512, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# 找到可以用于训练的 GPU\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Using {} device\".format(device))\n",
    "\n",
    "# 定义模型\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(3*224*224, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 损失函数与优化器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "定义一个损失函数和一个优化器。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义训练函数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在单个训练循环中，模型对训练数据集进行预测（分批提供给它），并反向传播预测误差从而调整模型的参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        # 计算预测误差\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # 反向传播\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 定义测试函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()\n",
    "    test_loss, correct = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 进行训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 2.374072  [    0/  600]\n",
      "loss: 0.342455  [  400/  600]\n",
      "Test Error: \n",
      " Accuracy: 89.7%, Avg loss: 0.334947 \n",
      "\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 0.634471  [    0/  600]\n",
      "loss: 0.167499  [  400/  600]\n",
      "Test Error: \n",
      " Accuracy: 93.2%, Avg loss: 0.198403 \n",
      "\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 0.303392  [    0/  600]\n",
      "loss: 0.028373  [  400/  600]\n",
      "Test Error: \n",
      " Accuracy: 98.3%, Avg loss: 0.099204 \n",
      "\n",
      "Epoch 4\n",
      "-------------------------------\n",
      "loss: 0.056484  [    0/  600]\n",
      "loss: 0.214862  [  400/  600]\n",
      "Test Error: \n",
      " Accuracy: 99.7%, Avg loss: 0.052564 \n",
      "\n",
      "Epoch 5\n",
      "-------------------------------\n",
      "loss: 0.093788  [    0/  600]\n",
      "loss: 0.054188  [  400/  600]\n",
      "Test Error: \n",
      " Accuracy: 99.8%, Avg loss: 0.033578 \n",
      "\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "epochs = 5\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train(train_loader, model, loss_fn, optimizer)\n",
    "    test(test_loader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
